package com.os.opencv.java.chapter13;

import org.opencv.core.*;
import org.opencv.highgui.HighGui;
import org.opencv.imgproc.Imgproc;
import org.opencv.video.Video;
import org.opencv.videoio.VideoCapture;

import java.util.ArrayList;
import java.util.Random;

public class OpticalFlowLK {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        //打开时视屏文件
        VideoCapture vc = new VideoCapture();
        vc.open("pics/2023-11-29_19-50-14.mp4");

        //视屏打开失败时的处理
        if(!vc.isOpened()){
            System.out.println("unable to load video!");
            System.exit(-1);
        }
        //定义颜色数组，用随机色跟踪不同对象
        Scalar[] colors = new Scalar[100];
        Random rd = new Random();
        for(int i=0; i<100; i++){
            int r = rd.nextInt(256);
            int g = rd.nextInt(256);
            int b = rd.nextInt(256);
            colors[i] = new Scalar(b, g, r);
        }

        //读取用于比较的原始画面并转换为灰度图
        Mat frame0 = new Mat();
        Mat gray0 = new Mat();
        vc.read(frame0);
        Imgproc.cvtColor(frame0, gray0, Imgproc.COLOR_BGR2GRAY);

        //进行角点检测
        MatOfPoint mop = new MatOfPoint();  //用于存储角点检测的结果
        Imgproc.goodFeaturesToTrack(gray0, mop, 100, 0.1, 10, new Mat(), 7, false, 0.04);

        //角点检测结果转换数据类型以便于处理
        MatOfPoint2f p0 = new MatOfPoint2f(mop.toArray());
        MatOfPoint2f p1 = new MatOfPoint2f();
        Mat mask = Mat.zeros(frame0.size(), frame0.type());

        //与原始画面一样的mask
        while(true){
            //读取与袁术画面比较的新画面frame
            Mat frame = new Mat();
            vc.read(frame);

            //如果视屏文件读完，则退出循环
            if(frame.empty()){
                break;
            }

            //将frame转换为灰度图
            Mat gray = new Mat();
            Imgproc.cvtColor(frame, gray, Imgproc.COLOR_BGR2GRAY);

            //定义光流估计用参数
            MatOfByte status = new MatOfByte();
            MatOfFloat err = new MatOfFloat();
            TermCriteria criteria = new TermCriteria(TermCriteria.COUNT + TermCriteria.EPS, 30, 0.03);

            //光流估计函数，p0为用于跟踪的特征点，有3个返回值
            //p1在新画面中的p0，status0和1与err
            Video.calcOpticalFlowPyrLK(gray0, gray, p0, p1, status, err, new Size(31,31), 2, criteria);

            //数据类型转换，便于处理
            byte statusArray[] = status.toArray();
            Point p0Arr[] = p0.toArray();
            Point p1Arr[] = p1.toArray();

            //将跟踪成功的额特征点轨迹绘制出来
            ArrayList<Point> goodOnes = new ArrayList<>();

            for(int i=0; i<statusArray.length; i++){
                if(statusArray[i] == 1){
                    //仅status=1的游泳，如为0则表示特征点跟踪丢了
                    goodOnes.add(p1Arr[i]);
                    Imgproc.line(mask, p1Arr[i], p0Arr[i], colors[i], 2);
                    Imgproc.circle(frame, p1Arr[i], 5, colors[i], -1);
                }
            }

            //将跟踪结果在屏幕上显示
            Mat img = new Mat();
            Core.add(frame, mask, img);  //与mask进行加运算
            HighGui.imshow("frame", img);
            int keyBoard = HighGui.waitKey(150);

            //按esc键则退出
            if(keyBoard == 27) break;

            //更新原始画面和相关变量
            gray0 = gray.clone();
            Point[] goodOnesArr = new Point[goodOnes.size()];
            goodOnesArr = goodOnes.toArray(goodOnesArr);
            p0 = new MatOfPoint2f(goodOnesArr);
        }

        HighGui.waitKey(0);
    }
}
